package chapter04.udtf;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;

import java.util.ArrayList;
import java.util.List;

public class MyUDTF extends GenericUDTF {
    private List<String> outList = new ArrayList<>();

    @Override
    public StructObjectInspector initialize(StructObjectInspector argOIs) throws UDFArgumentException {
        // define output data column name and output type
        List<String> fieldName = new ArrayList<>();
        List<ObjectInspector> fieldOIs = new ArrayList<>();

        // add column name and type of input data
        fieldName.add("words");
        fieldOIs.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);

        return ObjectInspectorFactory.getStandardStructObjectInspector(fieldName, fieldOIs);
    }

    @Override
    public void process(Object[] objects) throws HiveException {
        // get original data
        String string = objects[0].toString();
        String separatorRegex = objects[1].toString();

        // split string by separator regex
        String[] fields = string.split(separatorRegex);

        // each result to write
        for (String field : fields) {
            // clear
            outList.clear();
            // add
            outList.add(field);
            //write
            forward(outList);
        }
    }

    @Override
    public void close() throws HiveException {
        outList.clear();
    }
}
